# coding=utf-8
# @author:      ChengJing
# @name:        create_data.py
# @datetime:    2021/6/29 10:50
# @software:    PyCharm
# @description:

import numpy as np
import pandas as pd
import wntr
from multiprocessing import Pool


class CreateData:
    """
        生成数据集
        1. normal_pressure 获取生成正常压力数据
        2. pipe_burst 获取爆管压力
        3. sensitive 获取高日高时（或者任意时刻）的灵敏度矩阵
    """
    def __init__(
            self,
            wn,
            burst_level,
            duration=None,
            burst_interval=1800,
            discharge_coeff=0.41):
        """
        Args:
            wn: epanet水力模型
            burst_level: list, 爆管程度
            duration: int，单位：秒（s），模型的爆管模拟总时长，默认值和水力模型一致；如果指定时间则按照指定时间模拟
            burst_interval: 爆管的时间间隔，单位：秒
            discharge_coeff: 爆管公式的系数
        """
        self.wn = wn
        if duration:
            self.duration = duration
        else:
            self.duration = self.wn.options.time.duration
        self.burst_level = np.array(burst_level)
        self.burst_interval = burst_interval
        self.discharge_coeff = discharge_coeff

    def normal_data(
            self,
            is_save=True,
            Pfname='./datas/normal_pressure.csv',
            Qfname='./datas/normal_flow.csv'):
        """
        Args:
            is_save: 是否保存数据
            Pfname: 保存正常压力的文件
            Qfname: 保存正常流量的文件
        Return:
            pdata: 正常压力数据，dataframe，(行：时刻；列：节点名称)
            qdata: 正常流量数据，dataframe，(行：时刻；列：水库名称)
        """
        sim = wntr.sim.WNTRSimulator(self.wn)
        result = sim.run_sim()
        wn.reset_initial_values()
        pdata = result.node['pressure'].iloc[:-1, :-3]
        qdata = -result.node['demand'].iloc[:-1, -3:]
        if is_save:
            pdata.to_csv(Pfname)
            qdata.to_csv(Qfname)
        return pdata, qdata

    def _get_diameter(self, node):
        """
        Args:
            node: 节点的名字，str
        Returns:
            diameter：节点node相连的最大的管径，float
        """
        diameter = []
        links = self.wn.get_links_for_node(node)
        for link in links:
            diameter.append(self.wn.get_link(link).diameter)
        return max(diameter)

    def _burst(self, node, area, start_time):
        """
        Args:
            node: str, 节点id
            area: float，爆管面积
            start_time: 时间，单位：秒，爆管开始的时间
        Returns:
            p_data: 节点node爆管之后的压力数据
            q_data: 节点node爆管之后的漏损流量
            Q: 爆管之后的管网流量
        """
        id = node
        # self.wn.options.time.duration = self.burst_interval * 2
        self.wn.options.time.duration = start_time + self.burst_interval
        # self.wn.options.time.pattern_start = start_time - self.burst_interval
        node = self.wn.get_node(node)
        node.remove_leak(self.wn)
        node.add_leak(self.wn,
                      area=area,
                      discharge_coeff=self.discharge_coeff,
                      start_time=start_time,
                      end_time=start_time + self.burst_interval
                      )
        sim = wntr.sim.WNTRSimulator(self.wn)
        result = sim.run_sim()
        self.wn.reset_initial_values()
        p_data = result.node['pressure'].iloc[-61:-1, :-3]
        q_data = result.node['leak_demand'][id].iloc[-61:-1]
        Q = -result.node['demand'].iloc[-61:-1, -3:]
        return p_data, q_data, Q

    def pipe_burst(self, start_time=1800, nodes=None):
        """
        Args:
            start_time: 时间，单位：秒，爆管开始的时间
            nodes: list，模拟爆管的节点id集合，如果不指定则默认所有节点均发生爆管
        """
        if nodes:
            node_list = nodes
        else:
            node_list = self.wn.junction_name_list
        for node in node_list:
            areas = np.pi * (self._get_diameter(node)**2) / \
                4 * self.burst_level
            for n, area in enumerate(areas):
                start = start_time
                P_datas = pd.DataFrame()
                Q_datas = pd.DataFrame()
                Q_all = pd.DataFrame()
                while start <= (self.duration - self.burst_interval):
                    p_data, q_data, q = self._burst(node, area, start)
                    P_datas = pd.concat([P_datas, p_data])
                    Q_datas = pd.concat([Q_datas, q_data])
                    Q_all = pd.concat([Q_all, q])
                    start += self.burst_interval
                P_datas.to_csv(
                    './datas/burst/' + 'P_' + node + '_' + '%.2f' %
                    self.burst_level[n] + '.csv')
                Q_datas.to_csv(
                    './datas/burst/' + 'Q_' + node + '_' + '%.2f' %
                    self.burst_level[n] + '.csv')
                Q_all.to_csv(
                    './datas/burst/' + 'Q_all_' + node + '_' + '%.2f' %
                    self.burst_level[n] + '.csv')
                print(node + '_' + '%.2f' % self.burst_level[n] + ' is ok!')
            np.savetxt(
                './datas/burst/' + node + '_' + 'burst_area' + '.csv',
                areas,
                fmt='%.6f',
                delimiter=',')
        self.wn.reset_initial_values()

    def sensitive(self, bursts, start_time=1800):
        """
        Args:
            bursts: float, 爆管程度
            start_time: 时间，单位：秒，爆管开始的时间
        """
        try:
            normal = pd.read_csv(
                './datas/normal_pressure.csv',
                header=0,
                index_col=0)
        except BaseException:
            normal = self.normal_data(is_save=False)
        pressure = normal.loc[start_time].values
        node_list = wn.junction_name_list
        sen_matrix = np.zeros((0, self.wn.num_junctions))
        for n, node in enumerate(node_list):
            area = np.pi * (self._get_diameter(node) ** 2) / 4 * bursts
            start = start_time
            df, bq, _ = self._burst(node, area, start)
            data = df.values[int(df.shape[0] / 2), :].reshape(-1)
            sen = (pressure - data) / (pressure[n] - data[n])
            # sen = (pressure - data) / bq.values[int(bq.shape[0] / 2)]
            sen_matrix = np.concatenate(
                (sen_matrix, sen.reshape(1, -1)), axis=0)
        np.savetxt('./datas/sensitive_' + str(bursts) + '_' + str(start_time / 3600) + '.csv', sen_matrix, fmt='%.4f', delimiter=',')
        self.wn.reset_initial_values()


# 定义一个函数便于多进程生成数据，减少生成数据的时间
def mulcreatedata(wn, burst_level, nodes):
    """
    定义一个函数，方便使用多进程进行
    Args:
        wn: epanet水力模型
        burst_level: list, 爆管程度
        nodes: list，模拟爆管的节点id集合，如果不指定则默认所有节点均发生爆管
    """
    model = CreateData(wn, burst_level)
    model.pipe_burst(nodes=nodes)


if __name__ == '__main__':
    burst_level = np.arange(0.2, 1.6, 1.6)
    wn = wntr.network.WaterNetworkModel(r'./inp/tmodel24.inp')
    datamodel = CreateData(wn, burst_level)
    datamodel.normal_data()
    # datamodel.pipe_burst()
    # datamodel.sensitive(burst_level=1.0, start_time=12*60*3600)
    # nodes = wn.junction_name_list

    # 以进程池的方式采用多进程生成数据
    # pool = Pool(10)
    # for i in range(10):
    #     if i == 9:
    #         n = nodes[i * 47:-1]
    #     else:
    #         n = nodes[i * 47:i * 47 + 47]
    #     pool.apply_async(mulcreatedata, args=(wn, burst_level, n))
    #     print('进程%d启动成功,计算%d个节点' % (i, len(n)))
    # pool.close()
    # pool.join()
    # print('数据生成成功！')

